# -*- coding: utf-8 -*-
from scipy import *
from scipy import interpolate
from scipy import special
import os
import pickle

'''
module to construct the interpolation function from external data
The interpolation data contains the values for
theta 0 pi/2
phi 0 pi
The function is extended to larger range using the symmetry of the function.

For practical reason the a value is just calculated from 0 to 4. For larger A the following asymptotic rules are used:
a -> inf
e1 -> -a (*)
e2 -> -a
e3 -> a (*)
p1 -> 0 (*)
p2 -> 0
Since the system stays always in uni-stable state, only the (*) data is really needed.

We use the following arbitrary rule:
p1 = p1_last*a_last/a
e1 = (e1_last+a_last)*a_last/a-a
e3 = (e3_last-a_last)*a_last/a+a
which follow the same asymptotic behavior.

Also the interpolation function for thermal effect is also done.
The asymptotic behavior of such function is more complicated. See the mathematica file:
SW_thermal.nb
'''

defualt_dat = 'SW.dat'
thermal_dat = 'SW_thermal.dat'

a_last = loadtxt(defualt_dat)[:,0].max()

from import_math import stable_check

# Construct the tck data from pickle
file_check = True
filelist = ['e1.pickle','e2.pickle','e3.pickle','p1.pickle','p2.pickle']
# Check the interpolation files are generated correctly.
for fname in filelist:
    file_check = file_check and os.path.exists(fname)
if file_check is False:
    from import_math import const_interpolate
    const_interpolate(defualt_dat)

# have to use static functions.    
tck_dict = {}
for fname in filelist:
    with open(fname) as tck_data:
        tck = pickle.load(tck_data)
        func_name= fname.split('.')[0]
        tck_dict[func_name] = tck
        #def tmp_func(a, theta):
            #res = interpolate.bisplev(a, theta, tck_dict[func_name])
        #func_name= fname.split('.')[0]
        #func_dict[func_name] = tmp_func 

# Construct the thermal data.
thermal_raw = loadtxt(thermal_dat)
points = thermal_raw[:,0:3]
a_max = thermal_raw[:,0].max()
q_max = thermal_raw[:,2].max()
values = thermal_raw[:,3]
thermal_int = interpolate.LinearNDInterpolator(points, values)

def thermal_asymp(a,q):
    res = pi * (exp(a/q)*special.i0(a/q)-1)/(exp(2*a/q)-1)
    return res

def thermal_phi(a, theta, q):
    # handle theta symmetry
    if theta > pi/2:
        theta = pi - theta
    if a > a_max and q <= q_max:
        res = (thermal_int((a_max, theta, q)) - thermal_asymp(a_max, q))*a_max/a + thermal_asymp(a, q)
    elif a <= a_max and q > q_max:
        res = (thermal_int((a, theta, q_max)) - thermal_asymp(a, q_max))*q_max/q + thermal_asymp(a, q)
    elif a > a_max and q > q_max:
        res = (thermal_int((a_max,theta, q_max)) - thermal_asymp(a_max, q_max))*a_max/a*q_max/q + thermal_asymp(a, q)
    else:
        res = thermal_int((a,theta,q))
    return float(res)

def p1(a, theta):
    if a <= a_last:
        res = interpolate.bisplev(a, theta, tck_dict['p1'])
    else:
        p1_last = interpolate.bisplev(a_last, theta, tck_dict['p1'])
        res = p1_last*a_last/a
    return res
def p2(a, theta):
    res = interpolate.bisplev(a, theta, tck_dict['p2'])
    return res
def e1(a, theta):
    if a <= a_last:
        res = interpolate.bisplev(a, theta, tck_dict['e1'])
    else:
        e1_last = interpolate.bisplev(a_last, theta, tck_dict['e1'])
        res = (e1_last+a_last)*a_last/a-a
    return res
def e2(a, theta):
    res = interpolate.bisplev(a, theta, tck_dict['e2'])
    return res
def e3(a, theta):
    if a <= a_last:
        res = interpolate.bisplev(a, theta, tck_dict['e3'])
    else:
        e3_last = interpolate.bisplev(a_last, theta, tck_dict['e3'])
        (e3_last-a_last)*a_last/a+a
    return res

# extend the interpolation function using symmetry
def np1(a, theta):
    if theta > pi/2:
        theta = pi - theta
        out = 2*pi - p1(a, theta)
    else:
        out = p1(a, theta)
    return out

def np2(a, theta):
    if theta > pi/2:
        theta = pi - theta
        out = 2*pi - p2(a, theta)
    else:
        out = p2(a, theta)
    return out

def ne1(a, theta):
    if theta > pi/2:
        theta = pi - theta
        out = e1(a, theta)
    else:
        out = e1(a, theta)
    return out

def ne2(a, theta):
    if theta > pi/2:
        theta = pi - theta
        out = e2(a, theta)
    else:
        out = e2(a, theta)
    return out

def ne3(a, theta):
    if theta > pi/2:
        theta = pi - theta
        out = e3(a, theta)
    else:
        out = e3(a, theta)
    return out


if __name__=='__main__':
    print thermal_phi(0.3,0.3,0.5)